#![allow(unused_imports)]
#![allow(unused_variables)]
#![allow(unused_mut)]
#![allow(clippy::single_char_pattern)]

use std::collections::HashMap;
use std::sync::Arc ;
use chrono_tz::Tz;
use opensrv_mysql::*;
use arrow::array::{self, Array, ArrayRef, StringArray};
use arrow::datatypes::{Schema, SchemaRef, Field, DataType};
use arrow::record_batch::RecordBatch ;
use arrow::error::ArrowError as ErrorCode;
use regex::RegexSet;



const MYSQL_VERSION: &str = "8.0.26";
const THIS_VERSION: &str = "0.1.0" ;
//
#[derive(Clone)]
pub struct DataBlock (pub RecordBatch) ;
impl DataBlock {
    pub fn create(schema: SchemaRef, columns: Vec<ArrayRef>) -> Self {
        DataBlock(RecordBatch::try_new(schema, columns).unwrap())
    }
    pub fn empty() -> Self {
        DataBlock(RecordBatch::new_empty(Arc::new(Schema::empty())))
    }

    pub fn string_block(name: &str, val: &str) -> Self {
        DataBlock::create(
            Arc::new(Schema::new(vec![Field::new(
                name,
                DataType::Utf8,
                true
            )])),
            vec![Arc::new(StringArray::from(vec![val]))]
        )
    }
    #[inline]
    pub fn schema(&self) -> SchemaRef {
        self.0.schema()
    }

    #[inline]
    pub fn num_columns(&self) -> usize {
        self.0.columns().len()
    }

    #[inline]
    pub fn column(&self, index: usize) -> &ArrayRef {
        self.0.column(index)
    }
}

type LazyBlockFunc = fn(&str) -> Option<DataBlock>;

pub struct MySQLFederated {
    mysql_version: String,
    this_version: String,
}

impl MySQLFederated {
    pub fn create() -> Self {
        MySQLFederated {
            mysql_version: MYSQL_VERSION.to_string(),
            this_version: THIS_VERSION.to_string()
        }
    }

    // Build block for select @@variable.
    // Format:
    // |@@variable|
    // |value|
    fn select_variable_block(name: &str, value: &str) -> Option<DataBlock> {
        Some(DataBlock::string_block(name, value))
    }

    // Build block for show variable statement.
    // Format is:
    // |variable_name| Value|
    // | xx          | yy   |
    fn show_variables_block(name: &str, value: &str) -> Option<DataBlock> {
        Some(DataBlock::create(
            Arc::new(Schema::new(vec![
                Field::new("Variable_name", DataType::Utf8, true),
                Field::new("Value", DataType::Utf8, true),
            ])),
            vec![
                Arc::new(StringArray::from(vec![name])),
                Arc::new(StringArray::from(vec![value])),
            ],
        ))
    }

    fn block_match_rule(
        &self,
        query: &str,
        rules: Vec<(&str, Option<DataBlock>)>,
    ) -> Option<DataBlock> {
        let regex_rules = rules.iter().map(|x| x.0).collect::<Vec<_>>();
        let regex_set = RegexSet::new(&regex_rules).unwrap();
        let matches = regex_set.matches(query);
        for (index, (_regex, data_block)) in rules.iter().enumerate() {
            if matches.matched(index) {
                return match data_block {
                    None => Some(DataBlock::empty()),
                    Some(data_block) => Some(data_block.clone()),
                };
            }
        }
        None
    }

    fn lazy_block_match_rule(
        &self,
        query: &str,
        rules: Vec<(&str, LazyBlockFunc)>,
    ) -> Option<DataBlock> {
        let regex_rules = rules.iter().map(|x| x.0).collect::<Vec<_>>();
        let regex_set = RegexSet::new(&regex_rules).unwrap();
        let matches = regex_set.matches(query);
        for (index, (_regex, func)) in rules.iter().enumerate() {
            if matches.matched(index) {
                return match func(query) {
                    None => Some(DataBlock::empty()),
                    Some(data_block) => Some(data_block),
                };
            }
        }
        None
    }

    // SELECT @@aa, @@bb as cc, @dd...
    // Block is built by the variables.
    fn select_variable_data_block(query: &str) -> Option<DataBlock> {
        let mut default_map = HashMap::new();
        // DBeaver.
        default_map.insert("tx_isolation", "REPEATABLE-READ");
        default_map.insert("session.tx_isolation", "REPEATABLE-READ");
        default_map.insert("transaction_isolation", "REPEATABLE-READ");
        default_map.insert("session.transaction_isolation", "REPEATABLE-READ");
        default_map.insert("session.transaction_read_only", "0");
        default_map.insert("time_zone", "UTC");
        default_map.insert("system_time_zone", "UTC");
        // 128M
        default_map.insert("max_allowed_packet", "134217728");
        default_map.insert("interactive_timeout", "31536000");
        default_map.insert("wait_timeout", "31536000");
        default_map.insert("net_write_timeout", "31536000");

        let mut fields = vec![];
        let mut values = vec![];

        let query = query.to_lowercase();
        // select @@aa, @@bb, @@cc as yy, @@dd
        let mut vars: Vec<&str> = query.split("@@").collect();
        if vars.len() > 1 {
            vars.remove(0);
            for var in vars {
                let var = var.trim_end_matches(|c| c == ' ' || c == ',');
                let vars_as: Vec<&str> = var.split(" as ").collect();
                if vars_as.len() == 2 {
                    // @@cc as yy:
                    // var_as is 'yy' as the field name.
                    let var_as = vars_as[1];
                    fields.push(Field::new(var_as, DataType::Utf8, true));

                    // var is 'cc'.
                    let var = vars_as[0];
                    let value = default_map.get(var).unwrap_or(&"0").to_string();
                    values.push(Arc::new(StringArray::from(vec![value])) as ArrayRef);
                } else {
                    // @@aa
                    // var is 'aa'
                    fields.push(Field::new(&format!("@@{}", var), DataType::Utf8, true));

                    let value = default_map.get(var).unwrap_or(&"0").to_string();
                    values.push(Arc::new(StringArray::from(vec![value])) as ArrayRef);
                }
            }
        }
        Some(DataBlock::create(Arc::new(Schema::new(fields)), values))
    }

    // Check SELECT @@variable, @@variable
    fn federated_select_variable_check(&self, query: &str) -> Option<DataBlock> {
        let rules: Vec<(&str, LazyBlockFunc)> = vec![
            ("(?i)^(SELECT @@(.*))", Self::select_variable_data_block),
            (
                "(?i)^(/\\* mysql-connector-java(.*))",
                Self::select_variable_data_block,
            ),
        ];
        self.lazy_block_match_rule(query, rules)
    }

    // Check SHOW VARIABLES LIKE.
    fn federated_show_variables_check(&self, query: &str) -> Option<DataBlock> {
        let rules: Vec<(&str, Option<DataBlock>)> = vec![
            // sqlalchemy < 1.4.30
            (
                "(?i)^(SHOW VARIABLES LIKE 'sql_mode'(.*))",
                Self::show_variables_block("sql_mode",
                                           "ONLY_FULL_GROUP_BY STRICT_TRANS_TABLES NO_ZERO_IN_DATE NO_ZERO_DATE ERROR_FOR_DIVISION_BY_ZERO NO_ENGINE_SUBSTITUTION"),
            ),
            (
                "(?i)^(SHOW VARIABLES LIKE 'lower_case_table_names'(.*))",
                Self::show_variables_block("lower_case_table_names",
                                           "0"),
            ),
            ("(?i)^(show collation where(.*))", Self::show_variables_block("", "")),
            ("(?i)^(SHOW VARIABLES(.*))", Self::show_variables_block("", "")),
        ];
        self.block_match_rule(query, rules)
    }

    // Check for SET or others query, this is the final check of the federated query.
    fn federated_mixed_check(&self, query: &str) -> Option<DataBlock> {
        let rules: Vec<(&str, Option<DataBlock>)> = vec![
            (
                "(?i)^(SELECT VERSION())",
                Self::select_variable_block(
                    "version()",
                    format!("{}-{}", self.mysql_version, self.this_version.clone()).as_str(),
                ),
            ),
            // Txn.
            ("(?i)^(ROLLBACK(.*))", None),
            ("(?i)^(COMMIT(.*))", None),
            // Set.
            ("(?i)^(SET NAMES(.*))", None),
            ("(?i)^(SET character_set_results(.*))", None),
            ("(?i)^(SET FOREIGN_KEY_CHECKS(.*))", None),
            ("(?i)^(SET AUTOCOMMIT(.*))", None),
            ("(?i)^(SET sql_mode(.*))", None),
            ("(?i)^(SET @@(.*))", None),
            ("(?i)^(SET SESSION TRANSACTION ISOLATION LEVEL(.*))", None),
            // DBeaver.
            ("(?i)^(SHOW WARNINGS)", None),
            ("(?i)^(/\\* ApplicationName=(.*)SHOW WARNINGS)", None),
            ("(?i)^(/\\* ApplicationName=(.*)SHOW PLUGINS)", None),
            ("(?i)^(/\\* ApplicationName=(.*)SHOW COLLATION)", None),
            ("(?i)^(/\\* ApplicationName=(.*)SHOW CHARSET)", None),
            ("(?i)^(/\\* ApplicationName=(.*)SHOW ENGINES)", None),
            ("(?i)^(/\\* ApplicationName=(.*)SELECT @@(.*))", None),
            ("(?i)^(/\\* ApplicationName=(.*)SHOW @@(.*))", None),
            (
                "(?i)^(/\\* ApplicationName=(.*)SET net_write_timeout(.*))",
                None,
            ),
            (
                "(?i)^(/\\* ApplicationName=(.*)SET SQL_SELECT_LIMIT(.*))",
                None,
            ),
            ("(?i)^(/\\* ApplicationName=(.*)SHOW VARIABLES(.*))", None),
        ];

        self.block_match_rule(query, rules)
    }

    // Check the query is a federated or driver setup command.
    // Here we fake some values for the command which Databend not supported.
    pub fn check(&self, query: &str) -> Option<DataBlock> {
        // First to check the select @@variables.
        let select_variable = self.federated_select_variable_check(query);
        if select_variable.is_some() {
            return select_variable;
        }

        // Then to check the show variables like ''.
        let show_variables = self.federated_show_variables_check(query);
        if show_variables.is_some() {
            return show_variables;
        }

        // Last check.
        self.federated_mixed_check(query)
    }
}

pub struct DFQueryResultWriter<'a, W: std::io::Write> {
    inner: Option<QueryResultWriter<'a, W>>,
}

impl<'a, W: std::io::Write> DFQueryResultWriter<'a, W> {
    pub fn create(inner: QueryResultWriter<'a, W>) -> DFQueryResultWriter<'a, W> {
        DFQueryResultWriter::<'a, W> { inner: Some(inner) }
    }

    pub fn write(&mut self, query_result: anyhow::Result<(Vec<DataBlock>, String)>) -> anyhow::Result<()> {
        if let Some(writer) = self.inner.take() {
            match query_result {
                Ok((blocks, extra_info)) => Self::ok(blocks, extra_info, writer)?,
                Err(error) => Self::err(&error, writer)?,
            }
        }
        Ok(())
    }

    fn ok(
        blocks: Vec<DataBlock>,
        extra_info: String,
        dataset_writer: QueryResultWriter<'a, W>,
    ) -> anyhow::Result<()> {
        // XXX: num_columns == 0 may is error?
        let default_response = OkResponse {
            info: extra_info,
            ..Default::default()
        };

        if blocks.is_empty() || (blocks[0].num_columns() == 0) {
            dataset_writer.completed(default_response)?;
            return Ok(());
        }

        fn convert_field_type(field: &Field) -> anyhow::Result<ColumnType> {
            match field.data_type() {
                DataType::Int8 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::Int16 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::Int32 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::Int64 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::UInt8 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::UInt16 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::UInt32 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::UInt64 => Ok(ColumnType::MYSQL_TYPE_LONG),
                DataType::Float32 => Ok(ColumnType::MYSQL_TYPE_FLOAT),
                DataType::Float64 => Ok(ColumnType::MYSQL_TYPE_FLOAT),
                DataType::Utf8 => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
                DataType::Boolean => Ok(ColumnType::MYSQL_TYPE_SHORT),
                DataType::Date32 => Ok(ColumnType::MYSQL_TYPE_DATE),
                // DataType::DateTime32 => Ok(ColumnType::MYSQL_TYPE_DATETIME),
                // DataType::DateTime64 => Ok(ColumnType::MYSQL_TYPE_DATETIME),
                DataType::Null => Ok(ColumnType::MYSQL_TYPE_NULL),
                // DataType::Interval => Ok(ColumnType::MYSQL_TYPE_LONG),
                // DataType::Struct => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
                // DataType::Variant => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
                // DataType::VariantArray => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
                // DataType::VariantObject => Ok(ColumnType::MYSQL_TYPE_VARCHAR),
                _ => Err(anyhow::anyhow!(format!(
                    "Unsupported column type:{:?}",
                    field.data_type()
                ))),
            }
        }

        fn make_column_from_field(field: &Field) -> anyhow::Result<Column> {
            convert_field_type(field).map(|column_type| Column {
                table: "".to_string(),
                column: field.name().to_string(),
                coltype: column_type,
                colflags: ColumnFlags::empty(),
            })
        }

        fn convert_schema(schema: &SchemaRef) -> anyhow::Result<Vec<Column>> {
            schema.fields().iter().map(make_column_from_field).collect()
        }

        let block = blocks[0].clone();
        let utc: Tz = "UTC".parse().unwrap();

        match convert_schema(&block.schema()) {
            Err(error) => Self::err(&error, dataset_writer),
            Ok(columns) => {
                let columns_size = block.num_columns();
                let mut row_writer = dataset_writer.start(&columns)?;

                for block in &blocks {
                    let rows_size = block.column(0).len();
                    for row_index in 0..rows_size {
                        for col_index in 0..columns_size {
                            let block_schema = &block.schema() ;
                            let data_type = block_schema.fields()[col_index].data_type();
                            let data_array = block.column(col_index) ;
                            
                            if data_array.is_null(row_index) {
                                row_writer.write_col(None::<u8>)?;
                            }
                            match (data_type, data_array.data_type()) {
                                (DataType::Boolean, DataType::Boolean) => {
                                    let v = data_array.as_any().downcast_ref::<arrow::array::BooleanArray>().unwrap().value(row_index);
                                    row_writer.write_col(v as i8)?
                                }
/*
                                (DataType::Date16 ) => {
                                    row_writer.write_col(v.to_date(&utc).naive_local())?
                                }
                                (DataType::Date32, DataType::Date32) => {
                                    let v = data_array.as_any().downcast_ref::<arrow::array::Date32Array>().unwrap().value(row_index);
                                    row_writer.write_col(v.to_date(&utc).naive_local())?
                                }

                                (TypeID::DateTime32, DataValue::UInt64(v)) => {
                                    let data_type: &DateTime32Type =
                                        data_type.as_any().downcast_ref().unwrap();
                                    let tz = data_type.tz();
                                    let tz = tz.cloned().unwrap_or_else(|| "UTC".to_string());
                                    let tz: Tz = tz.parse().unwrap();
                                    row_writer.write_col(v.to_date_time(&tz).naive_local())?
                                }
                                (TypeID::DateTime64, DataValue::Int64(v)) => {
                                    let data_type: &DateTime64Type =
                                        data_type.as_any().downcast_ref().unwrap();
                                    let tz = data_type.tz();
                                    let tz = tz.cloned().unwrap_or_else(|| "UTC".to_string());
                                    let tz: Tz = tz.parse().unwrap();

                                    row_writer.write_col(
                                        v.to_date_time64(data_type.precision(), &tz)
                                            .naive_local()
                                            .format(data_type.format_string().as_str())
                                            .to_string(),
                                    )?
                                }
*/
                                (DataType::Utf8, DataType::Utf8) => {
                                    let v = data_array.as_any().downcast_ref::<arrow::array::StringArray>().unwrap().value(row_index);
                                    row_writer.write_col(v)?
                                }
/*
                                (DataType::Struct(_), DataType::Struct(_)) => {
                                    let v = data_array.as_any().downcast_ref::<arrow::array::StructArray>().unwrap().value(col_index);
                                    let serializer = data_type.create_serializer();
                                    row_writer.write_col(serializer.serialize_value(&val)?)?
                                }
                                (TypeID::Variant, DataValue::Json(_)) => {
                                    let serializer = data_type.create_serializer();
                                    row_writer.write_col(serializer.serialize_value(&val)?)?
                                }
                                (TypeID::VariantArray, DataValue::Json(_)) => {
                                    let serializer = data_type.create_serializer();
                                    row_writer.write_col(serializer.serialize_value(&val)?)?
                                }
                                (TypeID::VariantObject, DataValue::Json(_)) => {
                                    let serializer = data_type.create_serializer();
                                    row_writer.write_col(serializer.serialize_value(&val)?)?
                                }
*/
                                (DataType::Int64, DataType::Int64) => {
                                    let v = data_array.as_any().downcast_ref::<arrow::array::Int64Array>().unwrap().value(row_index);
                                    row_writer.write_col(v)?
                                }
                                (DataType::UInt64, DataType::UInt64) => {
                                    let v = data_array.as_any().downcast_ref::<arrow::array::UInt64Array>().unwrap().value(row_index);
                                    row_writer.write_col(v)?
                                }

                                (DataType::Float64, DataType::Float64) => {
                                    let v = data_array.as_any().downcast_ref::<arrow::array::Float64Array>().unwrap().value(row_index);
                                    row_writer.write_col(v)?
                                }
                                (_, v) => {
                                    return Err(anyhow::anyhow!(format!(
                                        "Unsupported column type:{:?}, expected type in schema: {:?}",
                                        v, data_type
                                    )));
                                }
                            }
                        }
                        row_writer.end_row()?;
                    }
                }
                row_writer.finish_with_info(&default_response.info)?;

                Ok(())
            }
        }
    }

    fn err(error: &anyhow::Error, writer: QueryResultWriter<'a, W>) -> anyhow::Result<()> {
/*
        if error.code() != ABORT_QUERY && error.code() != ABORT_SESSION {
            writer.error(ErrorKind::ER_UNKNOWN_ERROR, format!("{}", error).as_bytes())?;
        } else {
*/
            writer.error(
                ErrorKind::ER_ABORTING_CONNECTION,
                format!("{}", error).as_bytes(),
            )?;
/*
        }
*/

        Ok(())
    }
}

